import { onMounted, onUnmounted, watchEffect } from 'vue'
import '@tensorflow/tfjs'
import * as cocoSsd from '@tensorflow-models/coco-ssd'
import { ObjectDetection } from '@tensorflow-models/coco-ssd'

let model: ObjectDetection

cocoSsd
  .load({
    base: 'mobilenet_v2'
  })
  .then((res) => (model = res))

export function useCocoSsd(elRef) {
  const divs = []
  let mounted = false

  const handle = async () => {
    if (!elRef.value) return

    const isVideo = elRef.value instanceof HTMLVideoElement
    const c = document.createElement('canvas')
    const ctx = c.getContext('2d')

    async function predict() {
      if (!mounted) return

      // 媒体元素的父元素position需为relative
      const parent = elRef.value.parentElement
      let target = elRef.value

      if (isVideo) {
        // 视频尺寸缩放会影响坐标，转成图片
        const { width, height } = target.getBoundingClientRect()

        c.width = width
        c.height = height
        ctx.drawImage(target, 0, 0, width, height)

        target = c
      }

      let predictions = []
      try {
        predictions = await model.detect(target)
      } catch (e) {
        console.log(e)
      }

      for (let i = 0; i < divs.length; i++) {
        parent.removeChild(divs[i])
      }
      divs.splice(0)

      for (let n = 0; n < predictions.length; n++) {
        const { bbox, class: pclass, score } = predictions[n]
        if ((pclass == 'person' || pclass == 'car') && score > 0.66) {
          const title: any = document.createElement('div')
          title.innerText = `${pclass == 'person' ? '人' : pclass == 'car' ? '车' : ''}`

          title.style = `
            font-size: 12px;
            color: #fff;
            background: #09e;
            padding: 3px 10px;
            position: absolute;
            left: ${bbox[0]}px;
            top: ${bbox[1]}px;
          `

          const highlighter: any = document.createElement('div')
          highlighter.style = `
            border: 2px solid #0af;
            position: absolute;
            left: ${bbox[0]}px;
            top: ${bbox[1]}px;
            width: ${bbox[2]}px;
            height: ${bbox[3]}px;
          `

          parent.appendChild(highlighter)
          parent.appendChild(title)

          divs.push(highlighter)
          divs.push(title)
        }
      }

      if (isVideo) window.requestAnimationFrame(predict)
    }

    if (isVideo) elRef.value.addEventListener('loadeddata', predict)
    else elRef.value.addEventListener('load', predict)
  }

  watchEffect(() => {
    handle()
  })

  onMounted(() => {
    mounted = true
  })

  onUnmounted(() => {
    mounted = false
  })
}
